package com.luna.hbase.utils;

import cn.hutool.core.collection.CollUtil;
import com.luna.hbase.api.HbaseCallback;
import com.luna.hbase.api.HbaseOperations;
import com.luna.hbase.api.Mapper;
import com.luna.hbase.entity.HbaseCell;
import com.luna.hbase.entity.PutBatchModel;
import com.luna.hbase.entity.PutModel;
import com.luna.hbase.exception.HbaseException;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.filter.FirstKeyOnlyFilter;
import org.apache.hadoop.hbase.util.Bytes;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author Austin Wong
 * HBase Template
 * 2022/3/30 15:52
 * @since JDK1.8
 */
public class HbaseTemplate implements HbaseOperations {

    private final static Logger log = LoggerFactory.getLogger(HbaseTemplate.class);

    protected GenericObjectPool<Connection> pool;

    public HbaseTemplate(GenericObjectPool<Connection> pool) {
        this.pool = pool;
    }

    @Override
    public boolean createTable(String tableName, String... columnFamilies) {
        return this.execute(connection -> {
            Admin admin = connection.getAdmin();
            TableName name = TableName.valueOf(tableName);
            if (admin.tableExists(name)) {
                log.warn("table {} already exist!", name);
                return false;
            }

            // hbase2
            // List<ColumnFamilyDescriptor> list = Arrays.stream(columnFamilies)
            //         .map(columnFamily -> ColumnFamilyDescriptorBuilder.newBuilder(Bytes.toBytes(columnFamily)).build())
            //         .collect(Collectors.toList());
            // TableDescriptor tableDescriptor = TableDescriptorBuilder.newBuilder(name).setColumnFamilies(list).build();
            // admin.createTable(tableDescriptor);

            // hbase1.3
            HTableDescriptor tableDescriptor = new HTableDescriptor(name);
            for (String cf : columnFamilies) {
                HColumnDescriptor family = new HColumnDescriptor(cf);
                tableDescriptor.addFamily(family);
            }
            admin.createTable(tableDescriptor);
            log.info("create table {} successfully!", tableName);
            return true;

        });
    }

    @Override
    public List<String> listTableName() {
        return this.execute(connection -> {
            Admin admin = connection.getAdmin();
            TableName[] tableNames = admin.listTableNames();
            return Arrays.stream(tableNames).map(TableName::getNameAsString).collect(Collectors.toList());
        });
    }

    @Override
    @Deprecated
    public Map<String, List<HbaseCell>> scan(String tableName, Scan scan) {
        return this.execute(connection -> {
            Map<String, List<HbaseCell>> result;
            result = new HashMap<>(16);

            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                ResultScanner scanner = table.getScanner(scan);
                for (Result i : scanner) {
                    String rowKey = null;
                    List<HbaseCell> row = new ArrayList<>();
                    for (Cell cell : i.listCells()) {
                        if (Objects.isNull(rowKey)) {
                            rowKey = HBaseUtil.rowKey(cell);
                        }
                        row.add(new HbaseCell(rowKey, HBaseUtil.family(cell), HBaseUtil.column(cell), HBaseUtil.value(cell)));
                    }
                    result.put(rowKey, row);
                }
            } catch (IOException e) {
                log.error("An exception occurred while scanning table", e);
            }

            return result;
        });
    }

    @Override
    public boolean dropTable(String tableName) {
        return this.execute(connection -> {
            Admin admin = connection.getAdmin();
            TableName name = TableName.valueOf(tableName);
            try {
                if (admin.tableExists(name)) {
                    admin.disableTable(name);
                    admin.deleteTable(name);
                }
            } catch (IOException e) {
                e.printStackTrace();
                return false;
            }
            return true;
        });
    }

    @Override
    public long rowCount(String tableName) {
        return this.execute(connection -> {
            long rowCount = 0;
            ResultScanner resultScanner = null;
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Scan scan = new Scan();
                scan.setFilter(new FirstKeyOnlyFilter());
                resultScanner = table.getScanner(scan);
                for (Result result : resultScanner) {
                    rowCount += result.size();
                }
            } catch (IOException e) {
                log.error("An exception occurred when rowCount table", e);
            } finally {
                if (resultScanner != null) {
                    resultScanner.close();
                }
            }
            return rowCount;
        });
    }

    @Override
    public <T> List<T> list(String tableName, Scan scan, Mapper<T> mapper) {
        return this.execute(connection -> {
            List<T> list = new ArrayList<>();
            ResultScanner resultScanners = null;
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {

                int caching = scan.getCaching();
                // 如果caching未设置(默认是1)，将默认配置成5000
                if (caching == 1) {
                    scan.setCaching(5000);
                }

                resultScanners = table.getScanner(scan);
                for (Result result : resultScanners) {
                    list.add(mapper.mapping(result));
                }
            } catch (IOException e) {
                log.error("An exception occurred when list table", e);
            } finally {
                if (resultScanners != null) {
                    resultScanners.close();
                }
            }
            return list;
        });
    }

    @Override
    public void put(String tableName, String rowKey, String columnFamily, String[] columnQualifiers, String[] values) {
        this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Put put = new Put(Bytes.toBytes(rowKey));
                for (int i = 0, columnQualifiersLength = columnQualifiers.length; i < columnQualifiersLength; i++) {
                    String columnQualifier = columnQualifiers[i];
                    String value = values[i] == null ? "" : values[i];
                    put.addColumn(Bytes.toBytes(columnFamily), Bytes.toBytes(columnQualifier), Bytes.toBytes(value));
                }
                table.put(put);
            } catch (IOException e) {
                log.error("An exception occurred when put row into table", e);
            }
            return null;
        });
    }

    @Override
    public void putBatch(String tableName, String family, List<Map<String, String>> mapList) {
        this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                List<Put> puts = new ArrayList<>();
                for (Map<String, String> map : mapList) {
                    Put put = new Put(Bytes.toBytes(map.get("rowKey")));
                    for (Map.Entry<String, String> entry : map.entrySet()) {
                        if (!"rowKey".equals(entry.getKey())) {
                            String value = entry.getValue() == null ? "" : entry.getValue();
                            put.addColumn(Bytes.toBytes(family), Bytes.toBytes(entry.getKey()), Bytes.toBytes(value));
                        }
                    }
                    puts.add(put);
                }
                table.put(puts);
            } catch (IOException e) {
                log.error("An exception occurred when putBatch row into table", e);
            }
            return null;
        });
    }

    @Override
    public void putBatch(String tableName, List<PutBatchModel> batchModelList) {
        this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                batchModelList.forEach(b -> {
                    String familyName = b.getFamilyName();
                    List<PutModel> putModelList = b.getPutModelList();
                    List<Put> puts = new ArrayList<>();
                    putModelList.forEach(p -> {
                        if (CollUtil.isNotEmpty(p.getColumnValueMap().values())) {
                            Put put = new Put(Bytes.toBytes(p.getRowKey()));
                            for (Map.Entry<String, ?> entry : p.getColumnValueMap().entrySet()) {
                                Object value = entry.getValue() == null ? "" : entry.getValue();
                                put.addColumn(Bytes.toBytes(familyName), Bytes.toBytes(entry.getKey()), HBaseUtil.toByteArray(value));
                            }
                            puts.add(put);
                        }

                    });
                    try {
                        if (CollUtil.isNotEmpty(puts)) {
                            table.put(puts);
                        } else {
                            throw new HbaseException("所传数据为空");
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                });
            } catch (IOException e) {
                log.error("An exception occurred when putBatch row into table", e);
            }
            return null;
        });
    }

    @Override
    public void delete(String tableName, String rowKey) {
        this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Delete delete = new Delete(rowKey.getBytes());
                table.delete(delete);
            } catch (IOException e) {
                log.error("An exception occurred when delete row", e);
            }
            return null;
        });
    }

    @Override
    @Deprecated
    public List<HbaseCell> find(String tableName, String rowKey) {
        return this.execute(connection -> {
            List<HbaseCell> list = new ArrayList<>();
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Get get = new Get(rowKey.getBytes());
                Result result = table.get(get);
                for (Cell cell : result.rawCells()) {
                    String family = Bytes.toString(cell.getFamilyArray(), cell.getFamilyOffset(), cell.getFamilyLength());
                    String column = Bytes.toString(cell.getQualifierArray(), cell.getQualifierOffset(), cell.getQualifierLength());
                    String value = Bytes.toString(cell.getValueArray(), cell.getValueOffset(), cell.getValueLength());
                    list.add(new HbaseCell(rowKey, family, column, value));
                }
            } catch (IOException e) {
                log.error("An exception occurred when find row", e);
            }
            return list;
        });
    }

    @Override
    public <T> T get(String tableName, String rowKey, Mapper<T> mapper) {
        return this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Get get = new Get(rowKey.getBytes());
                Result result = table.get(get);
                return mapper.mapping(result);
            } catch (IOException e) {
                log.error("An exception occurred when get row", e);
            }
            return null;
        });
    }

    @Override
    public String get(String tableName, String rowKey, String family, String column) {
        return this.execute(connection -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName))) {
                Get get = new Get(rowKey.getBytes());
                get.addColumn(family.getBytes(), column.getBytes());
                Result result = table.get(get);
                Cell cell = result.listCells().stream().findFirst().orElse(null);
                table.close();
                if (Objects.nonNull(cell)) {
                    return HBaseUtil.value(cell);
                }
            } catch (IOException e) {
                log.error("An exception occurred when get value", e);
            }
            return null;
        });
    }

    @Override
    public <T> T execute(HbaseCallback<T> callback) {
        Assert.notNull(callback, "Callback of Hbase must not be null");

        Connection connection = null;
        try {
            connection = pool.borrowObject();
            return callback.doInHbase(connection);
        } catch (Exception e) {
            // e.printStackTrace();
            // return null;
            throw new HbaseException("An exception occurred when operate HBase via connection", e);
        } finally {
            if (Objects.nonNull(connection)) {
                pool.returnObject(connection);
            }
        }
    }


}
